package org.apache.solr.handler.component;

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

import java.net.MalformedURLException;
import java.util.Random;
import java.util.concurrent.*;

import org.apache.http.client.HttpClient;
import org.apache.solr.client.solrj.impl.HttpClientUtil;
import org.apache.solr.client.solrj.impl.LBHttpSolrServer;
import org.apache.solr.common.SolrException;
import org.apache.solr.common.params.ModifiableSolrParams;
import org.apache.solr.common.util.ExecutorUtil;
import org.apache.solr.common.util.NamedList;
import org.apache.solr.core.PluginInfo;
import org.apache.solr.util.DefaultSolrThreadFactory;
import org.apache.solr.util.plugin.PluginInfoInitialized;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class HttpShardHandlerFactory extends ShardHandlerFactory implements
		PluginInfoInitialized {
	protected static Logger log = LoggerFactory
			.getLogger(HttpShardHandlerFactory.class);

	// We want an executor that doesn't take up any resources if
	// it's not used, so it could be created statically for
	// the distributed search component if desired.
	//
	// Consider CallerRuns policy and a lower max threads to throttle
	// requests at some point (or should we simply return failure?)
	ThreadPoolExecutor commExecutor = new ThreadPoolExecutor(0,
			Integer.MAX_VALUE, 5, TimeUnit.SECONDS, // terminate idle threads
													// after 5 sec
			new SynchronousQueue<Runnable>(), // directly hand off tasks
			new DefaultSolrThreadFactory("httpShardExecutor"));

	private HttpClient defaultClient;
	LBHttpSolrServer loadbalancer;
	// default values:
	int soTimeout = 0;
	int connectionTimeout = 0;
	int maxConnectionsPerHost = 20;
	int corePoolSize = 0;
	int maximumPoolSize = Integer.MAX_VALUE;
	int keepAliveTime = 5;
	int queueSize = -1;
	boolean accessPolicy = false;

	public String scheme = "http://"; // current default values

	final Random r = new Random();

	// URL scheme to be used in distributed search.
	static final String INIT_URL_SCHEME = "urlScheme";

	// The core size of the threadpool servicing requests
	static final String INIT_CORE_POOL_SIZE = "corePoolSize";

	// The maximum size of the threadpool servicing requests
	static final String INIT_MAX_POOL_SIZE = "maximumPoolSize";

	// The amount of time idle threads persist for in the queue, before being
	// killed
	static final String MAX_THREAD_IDLE_TIME = "maxThreadIdleTime";

	// If the threadpool uses a backing queue, what is its maximum size (-1) to
	// use direct handoff
	static final String INIT_SIZE_OF_QUEUE = "sizeOfQueue";

	// Configure if the threadpool favours fairness over throughput
	static final String INIT_FAIRNESS_POLICY = "fairnessPolicy";

	/**
	 * Get {@link ShardHandler} that uses the default http client.
	 */
	public ShardHandler getShardHandler() {
		return getShardHandler(defaultClient);
	}

	/**
	 * Get {@link ShardHandler} that uses custom http client.
	 */
	public ShardHandler getShardHandler(final HttpClient httpClient) {
		return new HttpShardHandler(this, httpClient);
	}

	public void init(PluginInfo info) {
		NamedList args = info.initArgs;
		this.soTimeout = getParameter(args, HttpClientUtil.PROP_SO_TIMEOUT,
				soTimeout);
		this.scheme = getParameter(args, INIT_URL_SCHEME, "http://");
		this.scheme = (this.scheme.endsWith("://")) ? this.scheme : this.scheme
				+ "://";
		this.connectionTimeout = getParameter(args,
				HttpClientUtil.PROP_CONNECTION_TIMEOUT, connectionTimeout);
		this.maxConnectionsPerHost = getParameter(args,
				HttpClientUtil.PROP_MAX_CONNECTIONS_PER_HOST,
				maxConnectionsPerHost);
		this.corePoolSize = getParameter(args, INIT_CORE_POOL_SIZE,
				corePoolSize);
		this.maximumPoolSize = getParameter(args, INIT_MAX_POOL_SIZE,
				maximumPoolSize);
		this.keepAliveTime = getParameter(args, MAX_THREAD_IDLE_TIME,
				keepAliveTime);
		this.queueSize = getParameter(args, INIT_SIZE_OF_QUEUE, queueSize);
		this.accessPolicy = getParameter(args, INIT_FAIRNESS_POLICY,
				accessPolicy);

		BlockingQueue<Runnable> blockingQueue = (this.queueSize == -1) ? new SynchronousQueue<Runnable>(
				this.accessPolicy) : new ArrayBlockingQueue<Runnable>(
				this.queueSize, this.accessPolicy);

		this.commExecutor = new ThreadPoolExecutor(this.corePoolSize,
				this.maximumPoolSize, this.keepAliveTime, TimeUnit.SECONDS,
				blockingQueue,
				new DefaultSolrThreadFactory("httpShardExecutor"));

		ModifiableSolrParams clientParams = new ModifiableSolrParams();
		clientParams.set(HttpClientUtil.PROP_MAX_CONNECTIONS_PER_HOST,
				maxConnectionsPerHost);
		clientParams.set(HttpClientUtil.PROP_MAX_CONNECTIONS, 10000);
		clientParams.set(HttpClientUtil.PROP_SO_TIMEOUT, soTimeout);
		clientParams.set(HttpClientUtil.PROP_CONNECTION_TIMEOUT,
				connectionTimeout);
		clientParams.set(HttpClientUtil.PROP_USE_RETRY, false);
		this.defaultClient = HttpClientUtil.createClient(clientParams);

		try {
			loadbalancer = new LBHttpSolrServer(defaultClient);
		} catch (MalformedURLException e) {
			// should be impossible since we're not passing any URLs here
			throw new SolrException(SolrException.ErrorCode.SERVER_ERROR, e);
		}

	}

	private <T> T getParameter(NamedList initArgs, String configKey,
			T defaultValue) {
		T toReturn = defaultValue;
		if (initArgs != null) {
			T temp = (T) initArgs.get(configKey);
			toReturn = (temp != null) ? temp : defaultValue;
		}
		log.info("Setting {} to: {}", configKey, toReturn);
		return toReturn;
	}

	@Override
	public void close() {
		try {
			defaultClient.getConnectionManager().shutdown();
		} catch (Throwable e) {
			SolrException.log(log, e);
		}
		try {
			loadbalancer.shutdown();
		} catch (Throwable e) {
			SolrException.log(log, e);
		}
		try {
			ExecutorUtil.shutdownNowAndAwaitTermination(commExecutor);
		} catch (Throwable e) {
			SolrException.log(log, e);
		}
	}
}
